"""
PyG版本的随机图生成，与dgl.rand_graph函数兼容
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import numpy as np
from torch_geometric.data import Data
from ops.src.graph.graph import Graph


def rand_graph(num_nodes: int, num_edges: int) -> Graph:
    """
    生成一个随机图，与dgl.rand_graph函数兼容
    
    Args:
        num_nodes: 节点数量
        num_edges: 边数量
        
    Returns:
        随机生成的图
    """
    # 确保边数不超过可能的最大边数
    max_possible_edges = num_nodes * (num_nodes - 1)
    num_edges = min(num_edges, max_possible_edges)
    
    # 随机生成边索引
    edge_index = torch.zeros((2, num_edges), dtype=torch.long)
    
    # 生成所有可能的边
    all_edges = []
    for i in range(num_nodes):
        for j in range(num_nodes):
            if i != j:  # 排除自环
                all_edges.append((i, j))
    
    # 随机选择边
    selected_edges = np.random.choice(len(all_edges), num_edges, replace=False)
    for i, idx in enumerate(selected_edges):
        edge_index[0, i] = all_edges[idx][0]
        edge_index[1, i] = all_edges[idx][1]
    
    # 创建PyG图
    graph = Data(edge_index=edge_index, num_nodes=num_nodes)
    
    # 将结果转换为Graph类型，以保持与DGL接口的兼容性
    graph.__class__ = Graph
    
    return graph 